# https://gitee.com/yueyinqiu5990/tj12413601/blob/master/assignment2/question2/main_poloidal.py
import matplotlib.pyplot
import numpy
import scipy

import hydrogen
import my_special


def _calculate_and_draw(subplot, title, hydrogen_atom):
    theta = numpy.linspace(-numpy.pi, numpy.pi, 1000)
    phi = numpy.linspace(0, numpy.pi, 1000)
    theta, phi = numpy.meshgrid(theta, phi)

    p = hydrogen_atom.poloidal(theta, phi)
    r = numpy.abs(p) ** 2

    x = r * numpy.sin(phi) * numpy.cos(theta)
    y = r * numpy.sin(phi) * numpy.sin(theta)
    z = r * numpy.cos(phi)

    axes = matplotlib.pyplot.subplot(*subplot, projection="3d")
    matplotlib.pyplot.axis("off")
    matplotlib.pyplot.title(title)
    axes.plot_surface(x, y, z)


def main():
    h = hydrogen.Hydrogen(sph_harm_function=my_special.sph_harm)
    for m in range(0, 3):
        h.set_state(l_=2, m=m)
        _calculate_and_draw((2, 3, m + 1), f"$Y_{{2{m}}}$", h)

    h = hydrogen.Hydrogen(sph_harm_function=scipy.special.sph_harm)
    for m in range(0, 3):
        h.set_state(l_=2, m=m)
        _calculate_and_draw((2, 3, m + 1 + 3), "With SciPy", h)

    matplotlib.pyplot.subplots_adjust(
        left=0, bottom=0, right=1, top=0.9, wspace=0, hspace=0)
    matplotlib.pyplot.show()


if __name__ == "__main__":
    main()
